package com.elitesland.core.util;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.ObjectUtils;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;

/**
 * 树形数据转换工具类.
 *
 * @author Kaiser（wang shao）
 * @date 2020-11-17
 */
public class TreeDataUtil<T> {

    private final List<T> dataOriginal;
    private final Function<T, Object> getId;
    private final Function<T, Object> getParentId;
    private final BiConsumer<T, List<T>> setChildren;
    private Comparator<T> comparator;

    public TreeDataUtil(List<T> data, Function<T, Object> getId, Function<T, Object> getParentId) {
        this(data, getId, getParentId, null, null);
    }

    public TreeDataUtil(List<T> data, Function<T, Object> getId, Function<T, Object> getParentId, BiConsumer<T, List<T>> setChildren) {
        this(data, getId, getParentId, setChildren, null);
    }

    public TreeDataUtil(List<T> data, Function<T, Object> getId, Function<T, Object> getParentId, BiConsumer<T, List<T>> setChildren, Comparator<T> comparator) {
        this.dataOriginal = data;
        this.getId = getId;
        this.getParentId = getParentId;
        this.setChildren = setChildren;
        this.comparator = comparator;

        analyze();
    }

    public List<?> getRoots() {
        return isDefResult ? resultDef : resultCustom;
    }

    public Map<Object, ?> getChildrenMapping() {
        return isDefResult ? childrenMappingDef : childrenMappingCustom;
    }

    private TreeDataUtil<T> analyze() {
        if (CollectionUtils.isEmpty(dataOriginal)) {
            return this;
        }
        isDefResult = setChildren == null;

        if (isDefResult) {
            //使用默认的返回结果类型 TreeObj.TreeData
            analyzeForDef();
        } else {
            analyzeForCustom();
        }
        return this;
    }

    private void analyzeForDef() {
        int originalSize = (int) Math.ceil(dataOriginal.size() / 0.75);
        Set<Object> allIds = new HashSet<>(originalSize);
        //父子关系映射
        childrenMappingDef = new HashMap<>(originalSize);
        List<TreeData<T>> resultTemp = new ArrayList<>();

        //分析父子关系
        Object tempId, tempPId;
        TreeData<T> tempData;
        for (T t : dataOriginal) {
            tempId = getId.apply(t);
            tempPId = getParentId.apply(t);

            if (ObjectUtils.isEmpty(tempId)) {
                tempData = new TreeData<>(tempId, tempPId, t, new ArrayList<>(0));
            } else {
                tempData = new TreeData<>(tempId, tempPId, t, childrenMappingDef.computeIfAbsent(tempId, k -> new ArrayList<>(16)));
                allIds.add(tempId);
            }

            resultTemp.add(tempData);

            //归并至父节点下
            if (ObjectUtils.isNotEmpty(tempPId)) {
                childrenMappingDef.computeIfAbsent(tempPId, k -> new ArrayList<>(16)).add(tempData);
            }
        }

        //获取树的根节点
        if (CollectionUtils.isEmpty(allIds)) {
            resultDef = new ArrayList<>(0);
            return;
        }
        resultDef = new ArrayList<>(originalSize);

        for (TreeData<T> d : resultTemp) {
            tempPId = getParentId.apply(d.getData());

            if (ObjectUtils.isEmpty(tempPId) || !allIds.contains(tempPId)) {
                resultDef.add(d);
            }
        }

        // 排序
        if (comparator != null) {
            childrenMappingDef.values().forEach(t -> t.sort((o1, o2) -> comparator.compare(o1.getData(), o2.getData())));
            resultDef.sort((o1, o2) -> comparator.compare(o1.getData(), o2.getData()));
        }
    }

    private void analyzeForCustom() {
        int originalSize = (int) Math.ceil(dataOriginal.size() / 0.75);
        Set<Object> allIds = new HashSet<>(originalSize);
        //父子关系映射
        childrenMappingCustom = new HashMap<>(originalSize);

        Object tempId, tempPId;
        for (T t : dataOriginal) {
            tempId = getId.apply(t);
            tempPId = getParentId.apply(t);

            if (ObjectUtils.isEmpty(tempId)) {
                setChildren.accept(t, new ArrayList<>(0));
            } else {
                setChildren.accept(t, childrenMappingCustom.computeIfAbsent(tempId, k -> new ArrayList<>(0)));
                allIds.add(tempId);
            }

            if (ObjectUtils.isNotEmpty(tempPId)) {
                childrenMappingCustom.computeIfAbsent(tempPId, k -> new ArrayList<>(16)).add(t);
            }
        }


        if (CollectionUtils.isEmpty(allIds)) {
            resultCustom = new ArrayList<>(0);
            return;
        }

        resultCustom = new ArrayList<>(originalSize);
        for (T d : dataOriginal) {
            tempPId = getParentId.apply(d);

            if (ObjectUtils.isEmpty(tempPId) || !allIds.contains(tempPId)) {
                resultCustom.add(d);
            }
        }

        if (comparator != null) {
            childrenMappingCustom.values().forEach(t -> t.sort(comparator));
            resultCustom.sort(comparator);
        }
    }

    protected static class TreeData<T> {
        private Object id;
        private Object pId;
        private T data;
        private List<TreeData<T>> children;

        public TreeData(Object id, Object pId, T data, List<TreeData<T>> children) {
            this.id = id;
            this.pId = pId;
            this.data = data;
            this.children = children;
        }

        public TreeData() {
        }

        public Object getId() {
            return id;
        }

        public void setId(Object id) {
            this.id = id;
        }

        public Object getpId() {
            return pId;
        }

        public void setpId(Object pId) {
            this.pId = pId;
        }

        public T getData() {
            return data;
        }

        public void setData(T data) {
            this.data = data;
        }

        public List<TreeData<T>> getChildren() {
            return children;
        }

        public void setChildren(List<TreeData<T>> children) {
            this.children = children;
        }
    }

    /**
     * 是否是返回默认的树类型
     */
    private boolean isDefResult;

    /**
     * 父子关系映射
     */
    Map<Object, List<T>> childrenMappingCustom = null;
    Map<Object, List<TreeData<T>>> childrenMappingDef = null;

    /**
     * 返回结果
     */
    private List<T> resultCustom;
    private List<TreeData<T>> resultDef;
}
